"""
These are EXTRA functions used by the river code.

Parker MacCready
"""

def plot_rivers(riv_df, riv_df_raw, riv_list_to_plot):    
    """
    This plots river time series. Inputs 1 and 2 are DataFrames,
    and input 3 is a list of river names (a subset of those in the DataFrames).
    """
    
    import matplotlib.pyplot as plt   
    import numpy as np    
    import pandas as pd
       
    # set up the number of subplot panels
    nriv = len(riv_list_to_plot)
    ncol = int(np.ceil(np.sqrt(nriv)))
    nrow = int(np.ceil(float(nriv)/ncol))
        
    # set up the plots    
    plt.close()
    
    # shared time axis   
    t = riv_df.index
    tr = riv_df_raw.index
    
    count = 0
    for riv_name in riv_list_to_plot:
        q = riv_df[riv_name].values
        qr = riv_df_raw[riv_name].values
        
        irow = count/nrow
        icol = np.remainder(count,ncol)
        ax = plt.subplot(nrow, ncol, count + 1)
            
        # plot the data for one river in a panel
        ax.plot(t, q, 'ko--', label='Daily')
        # the raw values have NaN's so plot with just markers
        ax.plot(tr, qr, 'b.', markersize=2, label='Raw') # 2 appears to be the minimum    
        ax.set_title(riv_name.capitalize(), fontsize='small')
        # fix the time axis labeling
        ax.set_xticks(t)
        ax.set_xticklabels(t.day)
        if irow + 1 == nrow:
            ax.set_xlabel('Day of ' + str(t[0].month) + '/' + str(t[0].year))
        # fix the y-axis scaling
        (y0, y1) = ax.get_ylim()
        y1_new = y1 - np.mod(y1,100) + 100
        ax.set_ylim(0, y1_new)
        if count == 0:
            ax.legend(loc='best', fontsize='small')                  
        count = count + 1
        
    plt.show()
    